package yycore

import chisel3._
import chisel3.util._
import common.Constants._
import bus._
import chisel3.util.experimental.BoringUtils
import system.AddrSpace

object LSUOpType {
  // normal load/store
  // bit(1, 0) are size
  def lb   = "b0000".U
  def lh   = "b0001".U
  def lw   = "b0010".U
  def ld   = "b0011".U
  def lbu  = "b0100".U
  def lhu  = "b0101".U
  def lwu  = "b0110".U
  def sb   = "b1000".U
  def sh   = "b1001".U
  def sw   = "b1010".U
  def sd   = "b1011".U

  def isLoad(op: UInt): Bool = !op(3)
  def isStore(op: UInt): Bool = op(3)
  def MakeMask(addr: UInt, size: UInt) : UInt = {
    val wmask = Wire(Vec(8, Bool()))
    val numValid = 1.U << size
    wmask.zipWithIndex.map{ case (m, i) =>
      m := addr(2, 0) <= i.U && addr(2, 0) + numValid.asUInt() > i.U
    }
    wmask.asUInt()
  }
  def WriteData(addr: UInt, wdata: UInt): UInt = {
    val shiftwdata = Wire(UInt(XLEN.W))
    shiftwdata := MuxLookup(addr(2, 0), wdata, List(
      "b001".U  -> Cat(wdata(XLEN - 9, 0), Fill(8, "b0".U)),
      "b010".U  -> Cat(wdata(XLEN - 17, 0), Fill(16, "b0".U)),
      "b011".U  -> Cat(wdata(XLEN - 25, 0), Fill(24, "b0".U)),
      "b100".U  -> Cat(wdata(XLEN - 33, 0), Fill(32, "b0".U)),
      "b101".U  -> Cat(wdata(XLEN - 41, 0), Fill(40, "b0".U)),
      "b110".U  -> Cat(wdata(XLEN - 49, 0), Fill(48, "b0".U)),
      "b111".U  -> Cat(wdata(XLEN - 57, 0), Fill(56, "b0".U)),
    ))
    //shiftwdata := wdata << (addr(2, 0) << 3)
    shiftwdata
  }
  def genWdata(addr: UInt, wdata: UInt, size: UInt): UInt = {

    val partdata = MuxCase(wdata(XLEN - 1, 0), Array(
      (size === "b11".U)    ->  wdata(63, 0),
      (size === "b10".U)    ->  wdata(31, 0),
      (size === "b01".U)    ->  wdata(15, 0),
      (size === "b00".U)    ->  wdata(7, 0),
    )).asUInt()
    val shiftwdata = Wire(UInt(XLEN.W))
    shiftwdata := partdata << (addr(2, 0) << 3)
    shiftwdata
  }
}

class LSUBus extends Bundle {
  val addr = Output(UInt(AddrBits.W))
  val fun = Output(UInt(fuOpType.W))
  val rs2_data = Output(UInt(DataBits.W))

  def apply(addr: UInt, fun: UInt,/*en: Bool,*/ rs2: UInt) = {
    this.addr := addr
    this.fun := fun
    //this.mem_en := en
    this.rs2_data := rs2
    this
  }
}


class LSU extends Module {
  val io = IO(new Bundle() {
    val in = Flipped(Decoupled(new LSUBus))
    val out = Decoupled(UInt(XLEN.W))
    val dmem = new CoreLinkIO(XLEN)
    val mmio = new CoreLinkIO(XLEN)
    val isMMIO = Output(Bool())
  })


  val (valid, addr, fun, rs2_data) = (io.in.valid, io.in.bits.addr, io.in.bits.fun, io.in.bits.rs2_data)
  def access(valid: Bool, addr: UInt, fun: UInt, rs2_data: UInt): UInt = {
    this.valid := valid
    this.addr := addr
    this.fun := fun
    this.rs2_data := rs2_data
    io.out.bits
  }


  //
  // ===================== mem stage status =======================
  val load_data = WireInit(0.U(XLEN.W))
  val is_mmio = Wire(Bool())
  val shift = (addr(2, 0) << 3)
  val addrAtMMIO = AddrSpace.isMMIO(addr)
  is_mmio := valid && addrAtMMIO//(mem_reg.alu_out < 0x80000000L.U)
  when(valid && ((!is_mmio && io.dmem.resp.fire()) || (is_mmio && io.mmio.resp.fire()))) {
    load_data := Mux(is_mmio, io.mmio.resp.bits.rdata >> shift, io.dmem.resp.bits.rdata)
  }
  io.isMMIO := is_mmio

  // write data
  val size = fun(1, 0)
  val wdata = MuxCase(rs2_data(XLEN - 1, 0), Array(
    (size === "b11".U)    ->  rs2_data(63, 0),
    (size === "b10".U)    ->  rs2_data(31, 0),
    (size === "b01".U)    ->  rs2_data(15, 0),
    (size === "b00".U)    ->  rs2_data(7, 0),
  )).asUInt()
  //val mmio_data = (wdata << shift)


  // memory
  val memCmd = Mux(LSUOpType.isStore(fun), LinkBusCmd.write, LinkBusCmd.read)

  io.dmem.apply(
    valid = valid && !is_mmio,
    addr = Mux(valid && !is_mmio, addr, 0.U),
    wdata = Mux(valid && !is_mmio, LSUOpType.WriteData(addr, wdata), 0.U),
    size = Mux(valid && !is_mmio, size, 0.U),
    cmd = Mux(valid && !is_mmio, memCmd, 0.U))
  io.dmem.resp.ready := true.B

  // mmio
  io.mmio.apply(
    valid = valid && is_mmio,
    addr = Mux(valid && is_mmio, addr, 0.U),
    wdata = Mux(valid && is_mmio, LSUOpType.WriteData(addr, wdata), 0.U),
    size = Mux(valid && is_mmio, size, 0.U),
    cmd = Mux(valid && is_mmio, memCmd, 0.U))
  io.mmio.resp.ready := true.B

  // =================== write back data ============================
  val wb_ld_data = MuxCase(load_data, Array(
    (fun === LSUOpType.lb)   -> Cat(Fill(DataBits -  8, load_data(7)),   load_data(7, 0)),
    (fun === LSUOpType.lh)   -> Cat(Fill(DataBits - 16, load_data(15)),  load_data(15, 0)),
    (fun === LSUOpType.lw)   -> Cat(Fill(DataBits - 32, load_data(31)),  load_data(31, 0)),
    (fun === LSUOpType.lbu)  -> Cat(Fill(DataBits -  8, 0.U),            load_data(7, 0)),
    (fun === LSUOpType.lhu)  -> Cat(Fill(DataBits - 16, 0.U),            load_data(15, 0)),
    (fun === LSUOpType.lwu)  -> Cat(Fill(DataBits - 32, 0.U),            load_data(31, 0))
  ))

  io.in.ready := Mux(is_mmio, io.mmio.req.ready, io.dmem.req.ready)
  io.out.bits := wb_ld_data
  io.out.valid := Mux(!valid, true.B, Mux(is_mmio, io.mmio.resp.valid, io.dmem.resp.valid))

}
